import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor
from sklearn.metrics import accuracy_score, roc_curve, auc, mean_squared_error, r2_score
import matplotlib.pyplot as plt
from sklearn.preprocessing import LabelBinarizer
from imblearn.over_sampling import SMOTE
from Index.dataset import DataSet
import Index.dataset

if __name__ == '__main__':

    pd.set_option('display.max_rows', None)  # 显示所有行
    pd.set_option('display.max_columns', None)  # 显示所有列
    pd.set_option('display.width', None)  # 自动调整宽度，避免换行
    pd.set_option('display.max_colwidth', -1)  # 不限制列宽，显示完整内容
    np.set_printoptions(threshold=np.inf)

    # 对于模型种类进行选择
    modelType = 0
    # 加载示例数据集
    # df = DataSet('000001.SZ', '20200102', '20210102').data
    # df.loc[(df['Fmark'] == 2) | (df['Fmark'] == -2), 'Fmark'] = 0
    # #print(df['Fmark'])
    # for i in df.index:
    #     if df.loc[i, 'Fmark'] not in [1, -1, 0]: print('wwwwwww', i, df.loc[i, 'Fmark'])
    #
    # X = df.drop(columns=['trade_date', 'Fmark', 'close', 'high', 'low', 'EMA12', 'change', 'open'])
    # y = df['Fmark']
    # X = X[0:-1]
    # y = y[1:]
    # y = y.astype('int64')

    X, y = Index.dataset.StackedDataSet(['000001.SZ', '000006.SZ'], '20200201', '20210201')
    X = X[:, 1:] # 丢掉时间列
    for i in range(y.shape[0]):
        if y[i] == 2 or y[i] == -2:y[i] = 0
    X, y = SMOTE().fit_resample(X, y)

    # 将数据集划分为训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=123)

    if(modelType == 0):

        # 创建随机森林分类器对象
        rf_classifier = RandomForestClassifier(n_estimators=100, random_state=42)

        # 在训练集上训练模型
        rf_classifier.fit(X_train, y_train)

        # 在测试集上进行预测
        y_pred = rf_classifier.predict(X_test)

        # 计算预测精度
        accuracy = accuracy_score(y_test, y_pred)
        print(f'Accuracy: {accuracy:.2f}')

        y_prob = rf_classifier.predict_proba(X_test)

        # 计算每个类别的ROC曲线和AUC
        fpr = dict()
        tpr = dict()
        roc_auc = dict()

        lb = LabelBinarizer()
        y_test_bin = lb.fit_transform(y_test)

        for i in range(len(lb.classes_)):
            fpr[i], tpr[i], _ = roc_curve(y_test_bin[:, i], y_prob[:, i])
            roc_auc[i] = auc(fpr[i], tpr[i])

        # 绘制ROC曲线
        plt.figure(figsize=(8, 6))
        for i in range(len(lb.classes_)):
            plt.plot(fpr[i], tpr[i], label=f'ROC curve (class {lb.classes_[i]}) (AUC = {roc_auc[i]:.2f})')

        plt.plot([0, 1], [0, 1], 'k--', lw=2)
        plt.xlim([0.0, 1.0])
        plt.ylim([0.0, 1.05])
        plt.xlabel('False Positive Rate')
        plt.ylabel('True Positive Rate')
        plt.title('301111.SZ.csv ROC Curve')
        plt.legend(loc='lower right')
        plt.show()

    elif(modelType == 1) :
        # 创建随机森林回归器对象
        rf_regressor = RandomForestRegressor(n_estimators=100, random_state=42)

        # 在训练集上训练模型
        rf_regressor.fit(X_train, y_train)

        # 在测试集上进行预测
        y_pred = rf_regressor.predict(X_test)

        # 计算预测性能指标（如均方误差和R²分数）
        mse = mean_squared_error(y_test, y_pred)
        r2 = r2_score(y_test, y_pred)
        print(f'Mean Squared Error: {mse:.2f}')
        print(f'R² Score: {r2:.2f}')

        # 准备新数据（例如，从同一数据源获取的最新数据）
        # 这里仅作为示例，假设新数据是 X 中的前几行
        new_data = X.head()

        # 使用模型进行预测
        predicted_values = rf_regressor.predict(new_data)
        print("\nPredicted values:")
        print(predicted_values)